library(data.table)
library(gridExtra)
library(grid)
library(dplyr)
library(tidyr)
library(ggplot2)
library(stringr)
library(readr)
library(latex2exp)

## Paper or presentation
paper <- TRUE

## data directory
dir.project <- './'
dir.data <- paste0(dir.project, 'data/')
dir.raw <- paste0(dir.data, 'raw/')
dir.generated <- paste0(dir.data, 'generated/')

## figures directory for output
dir.fig.paper <-
  paste0(dir.project, 'doc/paper/figures/simulation/')
dir.fig.presentation <-
  paste0(dir.project, 'results/figures/presentation/')

source(paste0(dir.project, 'src/r/figures.R'))

fileName <- paste0(dir.project, 'calibration_version.txt')
calibration_version <-
  read_file(fileName) %>% str_replace_all("[\r\n]" , "")

#############
## FIGURES ##
#############

## effect of price drop given taxes and transfer

dat.equilibrium <-
  fread(
    paste0(
      dir.generated,
      'matlab/simulation/',
      calibration_version,
      '_simulation.csv'
    )
  ) %>% mutate(pricedrop = pricedrop * 100)
dat.equilibrium_high <-
  fread(
    paste0(
      dir.generated,
      'matlab/simulation/',
      calibration_version,
      '_simulation_high_tax.csv'
    )
  )  %>% mutate(pricedrop = pricedrop * 100)
dat.equilibrium_fix_capital <-
  fread(
    paste0(
      dir.generated,
      'matlab/simulation/',
      calibration_version,
      '_simulation_fix_capital.csv'
    )
  )  %>% mutate(pricedrop = pricedrop * 100)

dat.equilibrium$robot_tax <- 'status quo'
dat.equilibrium_high$robot_tax <- 'higher'
dat.combined <- rbind(dat.equilibrium, dat.equilibrium_high)

dat.taxes <-
  dat.combined %>% filter(robot_tax == 'status quo') %>%
  select(pricedrop,
         robot_tax,
         tax_rev_inc,
         exp_transfer,
         tax_rev_B,
         tax_rev_E,
         tax_rev_S) %>%
  gather(var, value,-robot_tax,-pricedrop)

ggplot(dat.taxes, aes(x = pricedrop, y = value)) + geom_line(aes(color =
                                                                   var, linetype = robot_tax))

dat.taxes <-
  dat.combined  %>% select(pricedrop,
                           robot_tax,
                           tax_rev_inc,
                           exp_transfer,
                           tax_rev_B,
                           tax_rev_E,
                           tax_rev_S) %>%
  gather(var, value,-robot_tax,-pricedrop)

ggplot(dat.taxes, aes(x = pricedrop, y = value)) + geom_line(aes(color =
                                                                   var, linetype = robot_tax))

dat.taxes.diff <-
  dat.taxes %>% group_by(var, pricedrop) %>%
  arrange(var, pricedrop, robot_tax) %>%
  transmute(diff = first(value) - last(value))

ggplot(dat.taxes.diff, aes(x = pricedrop, y = diff)) + geom_line(aes(color =
                                                                       var))

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = welfare)) +
  xlab('pct. drop in robot price') + ylab('')

to.pdf(print(p),
       "price_change_welfare.pdf",
       width = width.single,
       height = height.single)

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = K_B, linetype = '1')) +
  geom_line(aes(y = K_E, linetype = '2')) +
  geom_line(aes(y = K_S, linetype = '3')) + theme_bw() +
  theme(legend.position = "top") + xlab('pct. drop in robot price') + ylab('') +
  scale_linetype_manual(
    name = '',
    values = c('1' = 1, '2' = 2, '3' = 3),
    labels = c('robots', 'equipment', 'structures')
  )

to.pdf(print(p),
       "price_change_capital.pdf",
       width = width.single,
       height = height.single)

ggplot(dat.equilibrium  , aes(x = pricedrop)) +
  geom_line(aes(y = K_B), col = col.R) +
  geom_line(data = dat.equilibrium_high  ,
            aes(y = K_B),
            col = col.R,
            linetype = 2) +
  fig.theme

ggplot(dat.equilibrium  , aes(x = pricedrop)) +
  geom_line(aes(y = K_E), col = col.R) +
  geom_line(data = dat.equilibrium_high  ,
            aes(y = K_E),
            col = col.R,
            linetype = 2) +
  fig.theme

ggplot(dat.equilibrium  , aes(x = pricedrop)) +
  geom_line(aes(y = K_S), col = col.R) +
  geom_line(data = dat.equilibrium_high  ,
            aes(y = K_S),
            col = col.R,
            linetype = 2) +
  fig.theme

ggplot(dat.equilibrium  , aes(x = pricedrop)) +
  geom_line(aes(y = Y_M * L_M + Y_R * L_R + Y_C * L_C + Y_B_M * K_B + Y_E *
                  K_E + Y_S * K_S),
            col = col.R) +
  geom_line(
    data = dat.equilibrium_high  ,
    aes(y = Y_M * L_M + Y_R * L_R + Y_C * L_C + Y_B_M * K_B + Y_E * K_E + Y_S *
          K_S),
    col = col.R,
    linetype = 2
  ) +
  fig.theme

ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = Y_M * L_M + Y_R * L_R + Y_C * L_C + Y_B_M * K_B + Y_E *
                  K_E + Y_S * K_S),
            col = col.R) +
  geom_line(
    data = dat.equilibrium_high,
    aes(y = Y_M * L_M + Y_R * L_R + Y_C * L_C + Y_B_M * K_B + Y_E * K_E + Y_S *
          K_S),
    col = col.R,
    linetype = 2
  ) +
  fig.theme

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = avg_wage_M), col = col.M) +
  geom_line(aes(y = avg_wage_R), col = col.R) +
  geom_line(aes(y = avg_wage_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('') +
  theme.serif

to.pdf(print(p),
       "price_change_avg_wages.pdf",
       width = width.double,
       height = height.double)

avg_wage <-
  p + geom_line(data = dat.equilibrium_high ,
                aes(y = avg_wage_M),
                col = col.M,
                linetype = 2) +
  geom_line(data = dat.equilibrium_high ,
            aes(y = avg_wage_R),
            col = col.R,
            linetype = 2) +
  geom_line(data = dat.equilibrium_high ,
            aes(y = avg_wage_C),
            col = col.C,
            linetype = 2)

to.pdf(
  print(avg_wage),
  "price_change_avg_wages_high_tax.pdf",
  width = width.double,
  height = height.double
)

p <- ggplot(dat.equilibrium_fix_capital , aes(x = pricedrop)) +
  geom_line(aes(y = avg_wage_M), col = col.M) +
  geom_line(aes(y = avg_wage_R), col = col.R) +
  geom_line(aes(y = avg_wage_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('')

to.pdf(
  print(p),
  "price_change_avg_wages_fix_capital.pdf",
  width = width.double,
  height = height.double
)


p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = welfare_part_M), col = col.M) +
  geom_line(aes(y = welfare_part_R), col = col.R) +
  geom_line(aes(y = welfare_part_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('')

to.pdf(print(p),
       "price_change_welfare.pdf",
       width = width.double,
       height = height.double)

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = welfare_part_M / part_M), col = col.M) +
  geom_line(aes(y = welfare_part_R / part_R), col = col.R) +
  geom_line(aes(y = welfare_part_C / part_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('')

to.pdf(print(p),
       "price_change_welfare_per_cap.pdf",
       width = width.double,
       height = height.double)

p <-
  p + geom_line(
    data = dat.equilibrium_high ,
    aes(y = welfare_part_M / part_M),
    col = col.M,
    linetype = 2
  ) +
  geom_line(
    data = dat.equilibrium_high ,
    aes(y = welfare_part_R / part_R),
    col = col.R,
    linetype = 2
  ) +
  geom_line(
    data = dat.equilibrium_high ,
    aes(y = welfare_part_C / part_C),
    col = col.C,
    linetype = 2
  )

to.pdf(
  print(p),
  "price_change_welfare_per_cap_high_tax.pdf",
  width = width.double,
  height = height.double
)

p <- ggplot(dat.combined , aes(x = pricedrop)) +
  geom_line(aes(y = welfare, linetype = robot_tax)) +
  fig.theme  + xlab('pct. drop in robot price') + ylab('') +
  scale_linetype_manual(name = 'robot tax',
                        values = c('status quo' = 1, 'higher' = 2)) +
  theme(legend.position = "top")

to.pdf(print(p),
       "price_change_welfare_all.pdf",
       width = width.double,
       height = height.double)

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = Y_M), col = col.M) +
  geom_line(aes(y = Y_R), col = col.R) +
  geom_line(aes(y = Y_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('') +
  theme.serif

to.pdf(print(p),
       "price_change_Ys.pdf",
       width = width.double,
       height = height.double)

p <-
  p + geom_line(data = dat.equilibrium_high ,
                aes(y = Y_M),
                col = col.M,
                linetype = 2) +
  geom_line(data = dat.equilibrium_high ,
            aes(y = Y_R),
            col = col.R,
            linetype = 2) +
  geom_line(data = dat.equilibrium_high ,
            aes(y = Y_C),
            col = col.C,
            linetype = 2)

to.pdf(print(p),
       "price_change_Ys_high_tax.pdf",
       width = width.double,
       height = height.double)

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = emp_M), col = col.M) +
  geom_line(aes(y = emp_R), col = col.R) +
  geom_line(aes(y = emp_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('') +
  theme.serif

to.pdf(print(p),
       "price_change_emp.pdf",
       width = width.double,
       height = height.double)

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = part_M * 100), col = col.M) +
  geom_line(aes(y = part_R * 100), col = col.R) +
  geom_line(aes(y = part_C * 100), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('') +
  theme.serif

to.pdf(print(p),
       "price_change_part.pdf",
       width = width.double,
       height = height.double)


participation <-
  p + geom_line(data = dat.equilibrium_high ,
                aes(y = part_M * 100),
                col = col.M,
                linetype = 2) +
  geom_line(data = dat.equilibrium_high ,
            aes(y = part_R * 100),
            col = col.R,
            linetype = 2) +
  geom_line(data = dat.equilibrium_high ,
            aes(y = part_C * 100),
            col = col.C,
            linetype = 2)

to.pdf(
  print(participation),
  "price_change_part_high_tax.pdf",
  width = width.double,
  height = height.double
)

p <- ggplot(dat.equilibrium_fix_capital , aes(x = pricedrop)) +
  geom_line(aes(y = part_M), col = col.M) +
  geom_line(aes(y = part_R), col = col.R) +
  geom_line(aes(y = part_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('')

to.pdf(
  print(p),
  "price_change_part_fix_capital.pdf",
  width = width.double,
  height = height.double
)

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = pop_M), col = col.M) +
  geom_line(aes(y = pop_R), col = col.R) +
  geom_line(aes(y = pop_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('')

to.pdf(print(p),
       "price_change_pop.pdf",
       width = width.double,
       height = height.double)

p <- ggplot(dat.equilibrium, aes(x = pricedrop)) +
  geom_line(aes(y = part_M / pop_M), col = col.M) +
  geom_line(aes(y = part_R / pop_R), col = col.R) +
  geom_line(aes(y = part_C / pop_C), col = col.C) + fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('')

to.pdf(
  print(p),
  "price_change_part_conditional.pdf",
  width = width.double,
  height = height.double
)

robot_scale <- 1 / dat.equilibrium$K_B[[1]] * 0.36

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = K_B * robot_scale)) +
  fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('') +
  theme.serif

robots <-
  p + geom_line(data = dat.equilibrium_high , aes(y = K_B * robot_scale), linetype =
                  2)

to.pdf(
  print(robots),
  "price_change_robots_high_tax.pdf",
  width = width.double,
  height = height.double
)

gdp_normalize <- 1 / dat.equilibrium$Y[[1]]

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = Y * gdp_normalize)) +
  fig.theme + theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('') +
  theme.serif

output <-
  p + geom_line(data = dat.equilibrium_high , aes(y = Y * gdp_normalize), linetype =
                  2)

to.pdf(
  print(output),
  "price_change_output_high_tax.pdf",
  width = width.double,
  height = height.double
)

welfare_normalize <- 1 / dat.equilibrium$welfare[[1]]

p <- ggplot(dat.equilibrium , aes(x = pricedrop)) +
  geom_line(aes(y = welfare_normalize * welfare_part_M), col = col.M) +
  geom_line(aes(y = welfare_normalize * welfare_part_R), col = col.R) +
  geom_line(aes(y = welfare_normalize * welfare_part_C), col = col.C) +
  geom_line(aes(y = welfare_normalize * welfare_non_part), col = col.four) +
  geom_line(aes(y = welfare_normalize * welfare), col = col.five)  +
  theme(legend.position = "none") + xlab('pct. drop in robot price') + ylab('')

p <-
  p + geom_line(
    data = dat.equilibrium_high ,
    aes(y = welfare_normalize * welfare_part_M),
    col = col.M,
    linetype = 2
  ) +
  geom_line(
    data = dat.equilibrium_high ,
    aes(y = welfare_normalize * welfare_part_R),
    col = col.R,
    linetype = 2
  ) +
  geom_line(
    data = dat.equilibrium_high ,
    aes(y = welfare_normalize * welfare_part_C),
    col = col.C,
    linetype = 2
  ) +
  geom_line(
    data = dat.equilibrium_high ,
    aes(y = welfare_normalize * welfare_non_part),
    col = col.four,
    linetype = 2
  ) +
  geom_line(
    data = dat.equilibrium_high ,
    aes(y = welfare_normalize * welfare),
    col = col.five,
    linetype = 2
  )


to.pdf(
  print(p),
  "price_change_welfare_high_tax.pdf",
  width = width.single,
  height = height.single
)

# legends

leg_three_mrc <- ggplot(dat.equilibrium, aes(x = pricedrop)) +
  geom_line(aes(y = avg_wage_M, col = "a")) +
  geom_line(aes(y = avg_wage_R, col = "b")) +
  geom_line(aes(y = avg_wage_C, col = "c")) +
  xlab('pct. drop in robot price') + ylab('')  +
  theme_bw() + theme(legend.position = "bottom") +
  theme.serif +
  scale_color_manual(
    name = 'occupation',
    values = c('a' = col.M, 'b' = col.R, 'c' = col.C),
    labels = c('manual', 'routine', 'cognitive')
  )

to.pdf(
  fig.legend(leg_three_mrc),
  "legend_three_mrc.pdf",
  width = 4,
  height = 0.3
)

p <- ggplot(dat.equilibrium, aes(x = pricedrop)) +
  geom_line(aes(y = avg_wage_M, col = "a")) +
  geom_line(aes(y = avg_wage_R, col = "b")) +
  geom_line(aes(y = avg_wage_C, col = "c")) +
  geom_line(aes(y = avg_wage_C, col = "d")) +
  geom_line(aes(y = avg_wage_C, col = "e")) +
  xlab('pct. drop in robot price') + ylab('')  +
  theme_bw() + theme(legend.position = "bottom") +
  scale_color_manual(
    name = '',
    values = c(
      'a' = col.M,
      'b' = col.R,
      'c' = col.C,
      'd' = col.four,
      'e' = col.five
    ),
    labels = c('part. man.', 'part. rout.', 'part. cogn.', 'non-part.', 'total')
  )

to.pdf(fig.legend(p),
       "legend_groups_welfare.pdf",
       width = 4.5,
       height = 0.3)

legend.dat <- data.table(x_val = 1:4,
                         y_val = 5:8,
                         tax = as.factor(c(1, 1, 2, 2)))

low_high_tax_legend <-
  ggplot(legend.dat, aes(x = x_val, y = y_val)) + geom_line(aes(linetype =
                                                                  tax)) +
  scale_linetype_manual(
    name = unname(TeX('$t_B$')),
    values = (c(1, 2)),
    labels = c('1.55%', '15.5%')
  ) +
  theme_bw() +
  theme.serif +
  theme(legend.position = 'top')

to.pdf(
  fig.legend(low_high_tax_legend),
  "low_high_tax_legend.pdf",
  width = 2,
  height = 0.3
)

plots_top <-
  grid.arrange(
    robots + ggtitle('A: Robots per thsd. workers'),
    output + ggtitle('B: Output (normalized)'),
    ncol = 2,
    nrow = 1
  )

ggsave(
  paste0(dir.fig, "plots_simulate_given_tax_combined_top.pdf"),
  plot = plots_top,
  width = 10,
  height = 4.5,
  units = "cm"
)

plots_bottom <-
  grid.arrange(
    avg_wage + ggtitle('C: Avg. wages in USD'),
    participation + ggtitle('D: Empl.-to-pop. ratios'),
    ncol = 2,
    nrow = 1
  )

ggsave(
  paste0(dir.fig, "plots_simulate_given_tax_combined_bottom.pdf"),
  plot = plots_bottom,
  width = 10,
  height = 4.5,
  units = "cm"
)
